{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5cfa0417",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/agent/react_agent.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fb3144c-99b8-482d-a205-0f50877653aa",
   "metadata": {},
   "source": [
    "# ReAct Agent - A Simple Intro with Calculator Tools\n",
    "\n",
    "This is a notebook that showcases the ReAct agent over very simple calculator tools (no fancy RAG pipelines or API calls).\n",
    "\n",
    "We show how it can reason step-by-step over different tools to achieve the end goal."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "81d6fba5",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f2ea18c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8ac1778-0585-43c9-9dad-014d13d7460d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.agent import ReActAgent\n",
    "from llama_index.llms import OpenAI, ChatMessage\n",
    "from llama_index.tools import BaseTool, FunctionTool"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60ed8ac9-04fd-4deb-81a2-44cd8b4e2fd6",
   "metadata": {},
   "source": [
    "## Define Function Tools\n",
    "\n",
    "We setup some trivial `multiply` and `add` tools. Note that you can define arbitrary functions and pass it to the `FunctionTool` (which will process the docstring and parameter signature)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26472aaf-1a12-49f9-9fe6-cbf41dd15f88",
   "metadata": {},
   "outputs": [],
   "source": [
    "def multiply(a: int, b: int) -> int:\n",
    "    \"\"\"Multiply two integers and returns the result integer\"\"\"\n",
    "    return a * b\n",
    "\n",
    "\n",
    "multiply_tool = FunctionTool.from_defaults(fn=multiply)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df78ae85-bcf7-44c1-87ee-f301e646db20",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Add two integers and returns the result integer\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "add_tool = FunctionTool.from_defaults(fn=add)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b6c519f-6d74-4251-b38e-b8395a96d43b",
   "metadata": {},
   "source": [
    "## Run Some Queries\n",
    "\n",
    "### gpt-3.5-turbo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ab300f1-b054-46d9-b1c8-dbcd0d538e5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(model=\"gpt-3.5-turbo-instruct\")\n",
    "agent = ReActAgent.from_tools([multiply_tool, add_tool], llm=llm, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69bb1aa9-1ea3-4c88-a4f3-239b76392aa5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;5;200mThought: I need to use a tool to help me answer the question.\n",
      "assistant: Action: multiply\n",
      "assistant: Action Input: {\"a\": 2, \"b\": 4}\n",
      "Observation: 8\n",
      "assistant: Thought: I need to use a tool to help me answer the question.\n",
      "assistant: Action: add\n",
      "assistant: Action Input: {\"a\": 20, \"b\": 8}\n",
      "Observation: 28\n",
      "Thought: I can answer without using any more tools.\n",
      "Answer: 28\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "response = agent.chat(\"What is 20+(2*4)? Calculate step by step \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76112bb6-a291-4235-ad00-d4f6ff20adfe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 28"
     ]
    }
   ],
   "source": [
    "response_gen = agent.stream_chat(\"What is 20+2*4? Calculate step by step\")\n",
    "response_gen.print_response_stream()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d112393-ac6e-45d4-a58e-fec9745a79fc",
   "metadata": {},
   "source": [
    "### gpt-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bb7d49b-404c-4a46-9a84-1f7bb8792991",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(model=\"gpt-4\")\n",
    "agent = ReActAgent.from_tools([multiply_tool, add_tool], llm=llm, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "277f21d3-4f62-430e-b0fc-7561c64084d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;5;200mThought: I need to use the tools to help me answer the question. According to the order of operations in mathematics (BIDMAS/BODMAS), multiplication should be done before addition. So, I will first multiply 2 and 4, then add the result to 2.\n",
      "Action: multiply\n",
      "Action Input: {'a': 2, 'b': 4}\n",
      "\u001b[0m\u001b[1;3;34mObservation: 8\n",
      "\u001b[0m\u001b[1;3;38;5;200mThought: Now that I have the result of the multiplication, I need to add this to 2.\n",
      "Action: add\n",
      "Action Input: {'a': 2, 'b': 8}\n",
      "\u001b[0m\u001b[1;3;34mObservation: 10\n",
      "\u001b[0m\u001b[1;3;38;5;200mThought: I can answer without using any more tools.\n",
      "Answer: 10\n",
      "\u001b[0m10\n"
     ]
    }
   ],
   "source": [
    "response = agent.chat(\"What is 2+2*4\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6086fbe1-e82c-49ec-b410-cb951efab256",
   "metadata": {},
   "source": [
    "## View Prompts\n",
    "\n",
    "Let's take a look at the core system prompt powering the ReAct agent! \n",
    "\n",
    "Within the agent, the current conversation history is dumped below this line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "990597d5-f86c-44cb-ad91-24715a49a278",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(model=\"gpt-4\")\n",
    "agent = ReActAgent.from_tools([multiply_tool, add_tool], llm=llm, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad7964a5-a953-4a53-9865-6a0795cd2772",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prompt: agent_worker:system_prompt\n",
      "\n",
      "Value: \n",
      "You are designed to help with a variety of tasks, from answering questions     to providing summaries to other types of analyses.\n",
      "\n",
      "## Tools\n",
      "You have access to a wide variety of tools. You are responsible for using\n",
      "the tools in any sequence you deem appropriate to complete the task at hand.\n",
      "This may require breaking the task into subtasks and using different tools\n",
      "to complete each subtask.\n",
      "\n",
      "You have access to the following tools:\n",
      "{tool_desc}\n",
      "\n",
      "## Output Format\n",
      "To answer the question, please use the following format.\n",
      "\n",
      "```\n",
      "Thought: I need to use a tool to help me answer the question.\n",
      "Action: tool name (one of {tool_names}) if using a tool.\n",
      "Action Input: the input to the tool, in a JSON format representing the kwargs (e.g. {{\"input\": \"hello world\", \"num_beams\": 5}})\n",
      "```\n",
      "\n",
      "Please ALWAYS start with a Thought.\n",
      "\n",
      "Please use a valid JSON format for the Action Input. Do NOT do this {{'input': 'hello world', 'num_beams': 5}}.\n",
      "\n",
      "If this format is used, the user will respond in the following format:\n",
      "\n",
      "```\n",
      "Observation: tool response\n",
      "```\n",
      "\n",
      "You should keep repeating the above format until you have enough information\n",
      "to answer the question without using any more tools. At that point, you MUST respond\n",
      "in the one of the following two formats:\n",
      "\n",
      "```\n",
      "Thought: I can answer without using any more tools.\n",
      "Answer: [your answer here]\n",
      "```\n",
      "\n",
      "```\n",
      "Thought: I cannot answer the question with the provided tools.\n",
      "Answer: Sorry, I cannot answer your query.\n",
      "```\n",
      "\n",
      "## Current Conversation\n",
      "Below is the current conversation consisting of interleaving human and assistant messages.\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "prompt_dict = agent.get_prompts()\n",
    "for k, v in prompt_dict.items():\n",
    "    print(f\"Prompt: {k}\\n\\nValue: {v.template}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "502ffbf6-e1e9-46f6-8623-c0d8184fe274",
   "metadata": {},
   "source": [
    "### Customizing the Prompt\n",
    "\n",
    "For fun, let's try instructing the agent to output the answer along with reasoning in bullet points. See \"## Additional Rules\" section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0766d978-b011-40a6-bdce-b0ea566d2475",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.prompts import PromptTemplate\n",
    "\n",
    "react_system_header_str = \"\"\"\\\n",
    "\n",
    "You are designed to help with a variety of tasks, from answering questions \\\n",
    "    to providing summaries to other types of analyses.\n",
    "\n",
    "## Tools\n",
    "You have access to a wide variety of tools. You are responsible for using\n",
    "the tools in any sequence you deem appropriate to complete the task at hand.\n",
    "This may require breaking the task into subtasks and using different tools\n",
    "to complete each subtask.\n",
    "\n",
    "You have access to the following tools:\n",
    "{tool_desc}\n",
    "\n",
    "## Output Format\n",
    "To answer the question, please use the following format.\n",
    "\n",
    "```\n",
    "Thought: I need to use a tool to help me answer the question.\n",
    "Action: tool name (one of {tool_names}) if using a tool.\n",
    "Action Input: the input to the tool, in a JSON format representing the kwargs (e.g. {{\"input\": \"hello world\", \"num_beams\": 5}})\n",
    "```\n",
    "\n",
    "Please ALWAYS start with a Thought.\n",
    "\n",
    "Please use a valid JSON format for the Action Input. Do NOT do this {{'input': 'hello world', 'num_beams': 5}}.\n",
    "\n",
    "If this format is used, the user will respond in the following format:\n",
    "\n",
    "```\n",
    "Observation: tool response\n",
    "```\n",
    "\n",
    "You should keep repeating the above format until you have enough information\n",
    "to answer the question without using any more tools. At that point, you MUST respond\n",
    "in the one of the following two formats:\n",
    "\n",
    "```\n",
    "Thought: I can answer without using any more tools.\n",
    "Answer: [your answer here]\n",
    "```\n",
    "\n",
    "```\n",
    "Thought: I cannot answer the question with the provided tools.\n",
    "Answer: Sorry, I cannot answer your query.\n",
    "```\n",
    "\n",
    "## Additional Rules\n",
    "- The answer MUST contain a sequence of bullet points that explain how you arrived at the answer. This can include aspects of the previous conversation history.\n",
    "- You MUST obey the function signature of each tool. Do NOT pass in no arguments if the function expects arguments.\n",
    "\n",
    "## Current Conversation\n",
    "Below is the current conversation consisting of interleaving human and assistant messages.\n",
    "\n",
    "\"\"\"\n",
    "react_system_prompt = PromptTemplate(react_system_header_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67581223-f625-4b28-90aa-c8e5a232879e",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent.update_prompts({\"agent_worker:system_prompt\": react_system_prompt})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e1b2ce0-8f46-4d1d-8504-5b4b5d3f8478",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;5;200mThought: I need to use the add tool to help me answer the question.\n",
      "Action: add\n",
      "Action Input: {'a': 5, 'b': 3}\n",
      "\u001b[0m\u001b[1;3;34mObservation: 8\n",
      "\u001b[0m\u001b[1;3;38;5;200mThought: Now I need to add the result from the previous operation with 2.\n",
      "Action: add\n",
      "Action Input: {'a': 8, 'b': 2}\n",
      "\u001b[0m\u001b[1;3;34mObservation: 10\n",
      "\u001b[0m\u001b[1;3;38;5;200mThought: I can answer without using any more tools.\n",
      "Answer: The result of 5+3+2 is 10.\n",
      "- First, I added 5 and 3 using the add tool, which resulted in 8.\n",
      "- Then, I added the result (8) to 2 using the add tool, which resulted in 10.\n",
      "\u001b[0mThe result of 5+3+2 is 10.\n",
      "- First, I added 5 and 3 using the add tool, which resulted in 8.\n",
      "- Then, I added the result (8) to 2 using the add tool, which resulted in 10.\n"
     ]
    }
   ],
   "source": [
    "agent.reset()\n",
    "response = agent.chat(\"What is 5+3+2\")\n",
    "print(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v2",
   "language": "python",
   "name": "llama_index_v2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
